package com.bupt.ilink.service.bailian.impl;

import com.aliyun.bailian20231229.models.ListIndicesResponse;
import com.aliyun.bailian20231229.models.RetrieveResponse;
import com.bupt.ilink.config.BailianConfig;
import com.bupt.ilink.service.bailian.IndexService;
import com.bupt.ilink.tool.BailianFileTools;
import com.bupt.ilink.tool.BailianModelTools;
import com.bupt.ilink.utils.R;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;
@Service
public class IndexServiceImpl implements IndexService {

    @Autowired
    CreateIndexServiceImpl createIndexService;

    @Autowired
    ListCategoryServiceImpl listCategoryService;

    @Autowired
    SubmitIndexJobService submitIndexJobService;

    @Autowired
    GetIndexJobStatusService getIndexJobStatusService;

    @Autowired
    SubmitIndexAddDocumentsJobServiceImpl submitIndexAddDocumentsJobService;

    @Autowired
    ListFileService listFileService;

    @Autowired
    ListIndicesServiceImpl listIndicesService;

    @Autowired
    RetrieveServiceImpl retrieveService;

    String workspaceId = BailianConfig.workspaceId;

    @Override
    public R retrieveIndex(String query, String path) throws Exception {
        com.aliyun.bailian20231229.Client client = BailianConfig.createBailianClient();
        String className = BailianFileTools.splitPath(path).get(1);
        ListIndicesResponse listIndicesResponse = listIndicesService.listIndices(client, className, workspaceId);
        if (!BailianFileTools.checkResponse(listIndicesResponse)){
            return R.failure("获取索引列表失败", listIndicesResponse.getBody().getMessage());
        }
        String indexId = BailianFileTools.getIndexIdByName(listIndicesResponse, className);
        System.out.println(query);
        RetrieveResponse retrieve = retrieveService.retrieve(client, query, indexId, workspaceId);
        if (!BailianFileTools.checkResponse(retrieve)){
            return R.failure("检索接口调用失败", retrieve.getBody().getMessage());
        }
        List<Map<String, String>> output = new ArrayList<>();
        AtomicInteger i = new AtomicInteger(0);

        retrieve.getBody().getData().getNodes().forEach(node -> {
            System.out.println(node.getScore());
            System.out.println(node.getText());
            Map<String, String> metadata = (Map<String, String>) node.getMetadata();
            output.add(new HashMap<>());
            output.get(i.get()).put("参考内容", node.getText());
            System.out.println(i);
            output.get(i.getAndIncrement()).put("参考来源", metadata.get("doc_name"));
            System.out.println(i);
        });
        System.out.println(output);
        String prompt = BailianModelTools.buildPrompt(output.toString());
        String userInput = BailianModelTools.buildUserInput(query);
        String LLMOutput = BailianModelTools.runLLM(prompt, userInput);
        ObjectMapper mapper = new ObjectMapper();
        try {
            // 将JSON字符串解析为Map对象
            Map<String, String> responseMap = mapper.readValue(LLMOutput, new TypeReference<Map<String, String>>(){});

            // 提取“答案”和“参考来源”
            String answer = responseMap.get("答案");
            String reference = responseMap.get("参考来源");

            System.out.println("答案: " + answer);
            System.out.println("参考来源: " + reference);
            return R.ok("检索成功", responseMap);
        } catch (JsonProcessingException e) {
            e.printStackTrace();
            return R.failure("模型输出格式错误", LLMOutput);
        }
    }
}
